import os
from operator import itemgetter

REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:56379/0")
print(f"正在连接到Redis: {REDIS_URL}")

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_ollama import ChatOllama
from langchain_redis import RedisChatMessageHistory


# 获取session_id对应的ChatMessageHistory对象
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    return RedisChatMessageHistory(session_id, redis_url=REDIS_URL)



model = ChatOllama(
    model="llama3",
    temperature=0.5,
)

with_message_history = RunnableWithMessageHistory(
    itemgetter("input_messages") | model,
    get_session_history,
    input_messages_key="input_messages",
)

# 支持显示中文
while True:
    query = input('输入消息:')
    print('>>>>> 用户输入：:', query)
    if(query == 'exit'):
        print('>>>>> 用户退出，程序结束')
        break

    answer = with_message_history.invoke(
        input={'input_messages': query},
        config={'configurable': {'session_id': 'simpleGPT'}})
    print('>>>>> GPT:',answer.content)

